# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: taekwon

import experimental.prophet as prophet
from experimental.prophet.craps.universe import universe_all
from experimental.prophet.contrib.variable import last_valid


def single_product_features(symbol_info, feed):
  price_multiplier = symbol_info.price_multiplier or 1.0
  windows = ['1m', '2m', '3m', '4m', '5m', '10m', '20m', '30m', '60m']

  book = prophet.book(feed, inverse_product=symbol_info.inverse)
  trade = prophet.create_node('fastfeed.Trade', [feed.trade]).outputs
  # instrument = prophet.instrument(feed)

  timer_1s = prophet.timer('1s')
  ts = prophet.max(prophet.max(book.timestamp, trade.timestamp), timer_1s)
  with prophet.as_default_timestamp(ts):
    ask0_p = (book.ask0.price * price_multiplier).named_ref('b.ask0_p')
    ask0_q = book.ask0.qty.named_ref('b.ask0_q')
    bid0_p = (book.bid0.price * price_multiplier).named_ref('b.bid0_p')
    bid0_q = book.bid0.qty.named_ref('b.bid0_q')
    spread = (ask0_p - bid0_p).named_ref('b.spread')
    mid_p = ((ask0_p + bid0_p) / 2.).named_ref('b.mid_p')
    mid_p_std = prophet.time_moving_window(mid_p, windows).std()

    ask0_q_sum = prophet.time_moving_window(ask0_q, windows).sum()
    bid0_q_sum = prophet.time_moving_window(bid0_q, windows).sum()

    mid_p_ret = []
    for w in windows:
      v = (mid_p - prophet.time_shift(mid_p, w)).named_ref('mid_ret.' + w)
      mid_p_ret.append(v)

    book_vars = [
        ask0_p,
        ask0_q,
        *ask0_q_sum,
        bid0_p,
        bid0_q,
        *bid0_q_sum,
        spread,
        mid_p,
        *mid_p_std,
        *mid_p_ret
    ]

    inst_vars = []
    # try:
    #   oi = instrument.open_interest.named_ref('i.oi')
    #   oi_w = []
    #   for w in windows:
    #     oi_w.append((oi - timeseries.time_shift(oi, w)).named_ref('oi.' + w))
    #   inst_vars = [oi, *oi_w]
    # except Exception as e:
    #   print(nickname + ' has error in getting instrument due to ' + str(e) + '. Skipping...')

    tp = last_valid(trade.price * price_multiplier).named_ref('t.p')
    tp_max = prophet.time_moving_window(tp, windows).max()
    tp_min = prophet.time_moving_window(tp, windows).min()

    if not symbol_info.inverse:
      tq = trade.qty.named_ref('t.q')
      tv = (tp * tq).named_ref('t.v')
    else:
      tq = (trade.qty / tp).named_ref('t.q')
      tv = trade.qty.named_ref('t.v')
    tside = trade.side.named_ref('t.side')

    with prophet.control_if(tside == 1):
      buy_q_sum = prophet.time_moving_window(tq.named_ref('t.buy.q'), windows).sum()
      buy_v_sum = prophet.time_moving_window(tv.named_ref('t.buy.v'), windows).sum()

    with prophet.control_if(tside == 2):
      sell_q_sum = prophet.time_moving_window(tq.named_ref('t.sell.q'), windows).sum()
      sell_v_sum = prophet.time_moving_window(tv.named_ref('t.sell.v'), windows).sum()

    tq_sum = prophet.time_moving_window(tq, windows).sum()
    tv_sum = prophet.time_moving_window(tv, windows).sum()
    fil_tq_sum = []
    fil_tv_sum = []
    tvwaps = []
    vwap_ret = []
    for idx, window in enumerate(windows):
      tq = tq_sum[idx]
      fil_tq_sum.append(tq.named_ref(f't.q.sum{window}'))
      tv = tv_sum[idx]
      fil_tv_sum.append(tv.named_ref(f't.v.sum{window}'))
      vwap = (tv / tq).named_ref(f't.vwap.{window}')
      tvwaps.append(vwap)

      v = (vwap - prophet.time_shift(vwap, window)).named_ref('vwap_ret.' + window)
      vwap_ret.append(v)

    tp_std = prophet.time_moving_window(tp, windows).std()
    trade_vars = [
        tp,
        *tp_max,
        *tp_min,
        *tvwaps,
        *vwap_ret,
        *fil_tq_sum,
        *fil_tv_sum,
        *tp_std,
        *buy_q_sum,
        *buy_v_sum,
        *sell_q_sum,
        *sell_v_sum
    ]

    # prev_ask0_p = prophet.shift(ask0_p, 1)
    # prev_bid0_p = prophet.shift(bid0_p, 1)
    # ask_flip = (ask0_p > prev_ask0_p).named_ref('ask_flip')
    # bid_flip = (bid0_p < prev_bid0_p).named_ref('bid_flip')

    return book_vars + trade_vars + inst_vars  #,


def features(additional_symbols, date, is_y_feature=False):
  symbols = sorted(set(additional_symbols))
  universe = [s for s in universe_all(date) if s.nickname in symbols]
  vars = []

  for symbol_info in universe:
    feed = prophet.fastfeed_coin(symbol_info.product, symbol_info.sub_req)
    with prophet.namespace(symbol_info.nickname):
      if is_y_feature:
        vars += single_product_y_features(symbol_info, feed)
      else:
        vars += single_product_features(symbol_info, feed)
  return vars


def single_product_y_features(symbol_info, feed):
  windows = ['1m', '2m', '3m', '4m', '5m', '10m', '20m', '30m', '60m']
  book = prophet.book(feed, inverse_product=symbol_info.inverse)
  trade = prophet.create_node('fastfeed.Trade', [feed.trade]).outputs
  price_multiplier = symbol_info.price_multiplier or 1.0

  ask0_p = (book.ask0.price * price_multiplier)
  bid0_p = (book.bid0.price * price_multiplier)
  mid_p = ((ask0_p + bid0_p) / 2.)

  prev_y_mid_p_ret = []
  for w in windows:
    v = (mid_p - prophet.time_shift(mid_p, w)).named_ref('mid_ret.' + w)
    prev_mid_p_ret = prophet.shift(v, 1)
    prev_y_mid_p_ret.append(prophet.time_shift_y(prev_mid_p_ret, w))

  tp = last_valid(trade.price * price_multiplier)
  if not symbol_info.inverse:
    tq = trade.qty
    tv = (tp * tq)
  else:
    tq = (trade.qty / tp)
    tv = trade.qty
  tq_sum = prophet.time_moving_window(tq, windows).sum()
  tv_sum = prophet.time_moving_window(tv, windows).sum()

  prev_y_vwap_ret = []
  for idx, window in enumerate(windows):
    tq = tq_sum[idx]
    tv = tv_sum[idx]
    vwap = (tv / tq)
    v = (vwap - prophet.time_shift(vwap, window)).named_ref('vwap_ret.' + window)
    prev_vwap_ret = prophet.shift(v, 1)
    double_window = str(int(window[:-1]) * 2) + 'm'
    prev_y_vwap_ret.append(prophet.time_shift_y(prev_vwap_ret, double_window))

  return [*prev_y_mid_p_ret, *prev_y_vwap_ret]


# TODO(taekwon): Do this later
def samplers(target_coin, feature_prefix):
  graph = prophet.get_default_graph()

  if target_coin == 'btc':
    vol = graph.get_named_variable('%sokex-btcusd-quarter/trade_price.mmaxdiff30s' % feature_prefix)
    low_vol = (vol < 8.).named_ref('low_vol')  # 60%
    mid_vol = (8. <= vol and vol < 25.).named_ref('mid_vol')  # 60-95%
    high_vol = (25. <= vol).named_ref('high_vol')  # 95%

  elif target_coin == 'eth':
    vol = graph.get_named_variable('%sokex-ethusd-quarter/trade_price.mmaxdiff30s' % feature_prefix)
    low_vol = (vol < 0.15).named_ref('low_vol')  # 60%
    mid_vol = (0.15 <= vol and vol < 0.45).named_ref('mid_vol')  # 60-95%
    high_vol = (0.45 <= vol).named_ref('high_vol')  # 95%

  else:
    raise ValueError(target_coin)

  return [low_vol, mid_vol, high_vol]


# TODO(taekwon): Do this later
def trigger(target_coin, trigger_symbol_nick, feature_prefix):
  graph = prophet.get_default_graph()

  if target_coin == 'btc':
    trigger_nicks = set([trigger_symbol_nick, 'bitmex-xbtusd', 'okex-btcusd-quarter'])
  elif target_coin == 'eth':
    trigger_nicks = set([trigger_symbol_nick, 'bitmex-ethusd', 'okex-ethusd-quarter'])
  else:
    raise ValueError(target_coin)

  trigger = None
  for nickname in trigger_nicks:
    ask_flip = graph.get_named_variable('%s%s/ask_flip' % (feature_prefix, nickname))
    bid_flip = graph.get_named_variable('%s%s/bid_flip' % (feature_prefix, nickname))
    if trigger is None:
      trigger = ask_flip or bid_flip
    else:
      trigger = trigger or ask_flip or bid_flip

  return trigger.named_ref('trigger')
